

#include "stdio.h"
#include "tchar.h"
#include <userenv.h>
#include <windows.h>
#include <shellapi.h>
#include <iostream>
#include <wtsapi32.h>
#include <wchar.h>
using namespace std;
// #include <winuser.h>
#include "session.h"
#include "util.h"
#include <string>
#include <thread>


void Init();
BOOL IsInstalled();
BOOL Install();
BOOL Uninstall();
void LogEvent(LPCTSTR pszFormat, ...);
void WINAPI ServiceMain();
void WINAPI ServiceStrl(DWORD dwOpcode);
void DoWork();
void checkPrivilege();
bool Start(const std::wstring &strAppPath);
TCHAR szServiceName[] = _T("DhsReport");
BOOL bInstall;
SERVICE_STATUS_HANDLE hServiceStatus;
SERVICE_STATUS status;
DWORD dwThreadID;
PROCESS_INFORMATION processInfo = {0};
DWORD dwSessionId;
int APIENTRY WinMain(HINSTANCE hInstance,
                     HINSTANCE hPrevInstance,
                     LPSTR lpCmdLine,
                     int nCmdShow)
{

    // checkPrivilege();

    Init();

    dwThreadID = ::GetCurrentThreadId();

    SERVICE_TABLE_ENTRY st[] =
        {
            {szServiceName, (LPSERVICE_MAIN_FUNCTION)ServiceMain},
            {NULL, NULL}};

    if (_stricmp(lpCmdLine, "/install") == 0)
    {
        Install();
    }
    else if (_stricmp(lpCmdLine, "/uninstall") == 0)
    {
        Uninstall();
    }
    else
    {
        if (!::StartServiceCtrlDispatcher(st))
        {
            LogEvent(_T("Register Service Main Function Error!"));
        }
    }

    return 0;
}

void Init()
{
    hServiceStatus = NULL;
    status.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
    status.dwCurrentState = SERVICE_STOPPED;
    status.dwControlsAccepted = SERVICE_ACCEPT_STOP;
    status.dwWin32ExitCode = 0;
    status.dwServiceSpecificExitCode = 0;
    status.dwCheckPoint = 0;
    status.dwWaitHint = 0;
}


void WINAPI ServiceMain()
{
    // Register the control request handler
    status.dwCurrentState = SERVICE_START_PENDING;
    status.dwControlsAccepted = SERVICE_ACCEPT_STOP;

    // 注册服务控制
    hServiceStatus = RegisterServiceCtrlHandler(szServiceName, ServiceStrl);
    if (hServiceStatus == NULL)
    {
        LogEvent(_T("Handler not installed"));
        return;
    }
    SetServiceStatus(hServiceStatus, &status);

    status.dwWin32ExitCode = S_OK;
    status.dwCheckPoint = 0;
    status.dwWaitHint = 0;
    status.dwCurrentState = SERVICE_RUNNING;
    SetServiceStatus(hServiceStatus, &status);

    std::thread([]()
                { checkWindowSession(); })
        .detach();
    // 应用时将主要任务放于此
    while (status.dwCurrentState == SERVICE_RUNNING)
    {
        BEGIN_EXCEPT
        DoWork();
        END_EXCEPT
        Sleep(1000);
    }

    status.dwCurrentState = SERVICE_STOPPED;
    SetServiceStatus(hServiceStatus, &status);
    LogEvent(_T("Service stopped"));
}


void WINAPI ServiceStrl(DWORD dwOpcode)
{
    switch (dwOpcode)
    {
    case SERVICE_CONTROL_STOP:

        status.dwCurrentState = SERVICE_STOP_PENDING;
        SetServiceStatus(hServiceStatus, &status);

        TerminateProcess(processInfo.hProcess, 0);
        processInfo = {0};
        dwSessionId = 0;
        break;
    case SERVICE_CONTROL_PAUSE:
        break;
    case SERVICE_CONTROL_CONTINUE:
        LogEvent(L"==========%s", L"SERVICE_CONTROL_CONTINUE");

        break;
    case SERVICE_CONTROL_INTERROGATE:
        LogEvent(L"==========%s", L"SERVICE_CONTROL_INTERROGATE");

        break;
    case SERVICE_CONTROL_SHUTDOWN:
        LogEvent(L"==========%s", L"SERVICE_CONTROL_SHUTDOWN");

        break;
    default:
        LogEvent(_T("Bad service request"));
    }
}

BOOL IsInstalled()
{
    BOOL bResult = FALSE;

    // 打开服务控制管理器
    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);

    if (hSCM != NULL)
    {
        // 打开服务
        SC_HANDLE hService = ::OpenService(hSCM, szServiceName, SERVICE_QUERY_CONFIG);
        if (hService != NULL)
        {
            bResult = TRUE;
            ::CloseServiceHandle(hService);
        }
        ::CloseServiceHandle(hSCM);
    }
    return bResult;
}


BOOL Install()
{
    // if (IsInstalled())
    //     return TRUE;

        Uninstall();

      Sleep(2000);
    // 打开服务控制管理器
    SC_HANDLE hSCM = ::OpenSCManager(NULL, NULL, SC_MANAGER_ALL_ACCESS);
    if (hSCM == NULL)
    {
        MessageBox(NULL, _T("Couldn't open service manager"), szServiceName, MB_OK);
        return FALSE;
    }

    // Get the executable file path
    TCHAR szFilePath[MAX_PATH];
    ::GetModuleFileName(NULL, szFilePath, MAX_PATH);

    // 创建服务
    SC_HANDLE hService = ::CreateServiceW(
        hSCM, szServiceName, szServiceName,
        SERVICE_ALL_ACCESS, SERVICE_WIN32_OWN_PROCESS,
        SERVICE_AUTO_START, SERVICE_ERROR_NORMAL,
        szFilePath, NULL, NULL, _T(""), NULL, NULL);

    if (hService == NULL)
    {
        ::CloseServiceHandle(hSCM);
        MessageBox(NULL, _T("Couldn't create service"), szServiceName, MB_OK);
        return FALSE;
    }

    // SERVICE_DELAYED_AUTO_START_INFO info;
    // info.fDelayedAutostart = true;

    // ::ChangeServiceConfig2W(hService, SERVICE_CONFIG_DELAYED_AUTO_START_INFO, &info);

    StartServiceW(hService, 0, NULL);
    ::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);
    return TRUE;
}


BOOL Uninstall()
{
    if (!IsInstalled())
        return TRUE;

    SC_HANDLE hSCM = ::OpenSCManagerW(NULL, NULL, SC_MANAGER_ALL_ACCESS);

    if (hSCM == NULL)
    {
        // MessageBox(NULL, _T("Couldn't open service manager"), szServiceName, MB_OK);
        return FALSE;
    }

    SC_HANDLE hService = ::OpenServiceW(hSCM, szServiceName, SERVICE_ALL_ACCESS);

    if (hService == NULL)
    {
        ::CloseServiceHandle(hSCM);
        // MessageBox(NULL, _T("Couldn't open service"), szServiceName, MB_OK);
        return FALSE;
    }
      SERVICE_STATUS status;
    ::ControlService(hService, SERVICE_CONTROL_STOP, &status);

    // 删除服务
    BOOL bDelete = ::DeleteService(hService);
    ::CloseServiceHandle(hService);
    ::CloseServiceHandle(hSCM);

    if (bDelete)
        return TRUE;

    LogEvent(_T("Service could not be deleted"));
    return FALSE;
}


void LogEvent(LPCTSTR pFormat, ...)
{
    TCHAR chMsg[256];
    HANDLE hEventSource;
    LPTSTR lpszStrings[1];
    va_list pArg;

    va_start(pArg, pFormat);
    swprintf_s(chMsg, pFormat, pArg);
    va_end(pArg);

    lpszStrings[0] = chMsg;

    hEventSource = RegisterEventSource(NULL, szServiceName);
    if (hEventSource != NULL)
    {
        ReportEvent(hEventSource, EVENTLOG_INFORMATION_TYPE, 0, 0, NULL, 1, 0, (LPCTSTR *)&lpszStrings[0], NULL);
        DeregisterEventSource(hEventSource);
    }
}


void DoWork()
{

    using namespace std;

    char cWindowsDirectory[MAX_PATH];

    TCHAR szFilePath[MAX_PATH];
    ::GetModuleFileNameW(NULL, szFilePath, MAX_PATH);
    wchar_t *pos = wcsrchr(szFilePath, L'\\');
    if (pos != nullptr)
    {
        *pos = L'\0';
    }
    else
    {
    }
    wcscat_s(szFilePath, L"\\DhsReport.exe");

    Start(szFilePath);
}

bool Start(const std::wstring &strAppPath)
{

    HANDLE hToken = NULL;
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &hToken))
    {
        return false;
    }

    HANDLE hTokenDup = NULL;
    bool bRet = DuplicateTokenEx(hToken, TOKEN_ALL_ACCESS, NULL, SecurityIdentification, TokenPrimary, &hTokenDup);
    if (!bRet || hTokenDup == NULL)
    {
        CloseHandle(hToken);
        return false;
    }

    dwSessionId = WTSGetActiveConsoleSessionId();
    if (dwSessionId == 0)
    {
        return false;
    }
    // 把服务hToken的SessionId替换成当前活动的Session(即替换到可与用户交互的winsta0下)
    if (!SetTokenInformation(hTokenDup, TokenSessionId, &dwSessionId, sizeof(DWORD)))
    {

        LogEvent(_T("SetTokenInformation error"));
        DWORD nErr = GetLastError();
        CloseHandle(hTokenDup);
        CloseHandle(hToken);
        return false;
    }

    STARTUPINFO si;
    ZeroMemory(&si, sizeof(STARTUPINFO));

    si.cb = sizeof(STARTUPINFO);
    si.lpDesktop = _T("WinSta0\\Default");
    si.wShowWindow = SW_SHOW;
    si.dwFlags = STARTF_USESHOWWINDOW /*|STARTF_USESTDHANDLES*/;

    // 创建进程环境块
    LPVOID pEnv = NULL;
    bRet = CreateEnvironmentBlock(&pEnv, hTokenDup, FALSE);
    if (!bRet)
    {
        LogEvent(_T("CreateEnvironmentBlock error"));

        CloseHandle(hTokenDup);
        CloseHandle(hToken);
        return false;
    }

    if (pEnv == NULL)
    {
        CloseHandle(hTokenDup);
        CloseHandle(hToken);
        return false;
    }

    // 在活动的Session下创建进程
    ZeroMemory(&processInfo, sizeof(PROCESS_INFORMATION));
    DWORD dwCreationFlag = NORMAL_PRIORITY_CLASS | CREATE_NEW_CONSOLE | CREATE_UNICODE_ENVIRONMENT;

    if (!CreateProcessAsUserW(hTokenDup, NULL, (LPWSTR)strAppPath.c_str(), NULL, NULL, FALSE, dwCreationFlag, pEnv, NULL, &si, &processInfo))
    {
        LogEvent(_T("start DhsReport error"));
        DWORD nRet = GetLastError();
        CloseHandle(hTokenDup);
        CloseHandle(hToken);
        return false;
    }

    WaitForSingleObject(processInfo.hProcess, INFINITE); // 关闭句柄
    DestroyEnvironmentBlock(pEnv);
    CloseHandle(hTokenDup);
    CloseHandle(hToken);

    return true;
}